import numpy as np
import pandas as pd
from sklearn.preprocessing import LabelEncoder

# 设置随机种子（保证可复现）
np.random.seed(42)

# 生成样本量
n_samples = 1000

# 1. 年龄（18-64岁，右偏分布）
age = np.round(np.random.gamma(shape=6, scale=4, size=n_samples)).clip(18, 64)

# 2. 性别（二元分类，略微不平衡）
sex = np.random.choice(['male', 'female'], size=n_samples, p=[0.51, 0.49])

# 3. BMI（正态分布，均值26，标准差5）
bmi = np.random.normal(loc=26, scale=5, size=n_samples).clip(15, 45)

# 4. 子女数量（泊松分布，截断）
children = np.random.poisson(lam=0.8, size=n_samples).clip(0, 5)

# 5. 是否吸烟（与charges强相关）
smoker = np.random.choice(['yes', 'no'], size=n_samples, p=[0.2, 0.8])

# 6. 地区（均匀分布）
region = np.random.choice(['northeast', 'northwest', 'southeast', 'southwest'], size=n_samples)

# 7. 保险费用（charges，基于真实公式+噪声）
base_cost = 250 * age + 300 * (bmi - 25) + 500 * children
smoker_penalty = np.where(np.array(smoker) == 'yes', 20000, 0)
region_adjust = np.where(np.array(region) == 'southeast', 500,
                        np.where(np.array(region) == 'northeast', 300, 0))
noise = np.random.normal(loc=0, scale=1000, size=n_samples)
charges = base_cost + smoker_penalty + region_adjust + noise
charges = charges.clip(1000, 60000)  # 截断极端值

# 构建DataFrame
data = pd.DataFrame({
    'age': age,
    'sex': sex,
    'bmi': bmi,
    'children': children,
    'smoker': smoker,
    'region': region,
    'charges': charges
})

# 保存为CSV
data.to_csv('insurance_data.csv', index=False)